#include "Common.h"
#include "Mesh.h"
#include "Utils/Hash.h"
#include "Utils/KdTree.h"
#include "Graphics/OpenGL/OpenGLRenderer.h"
#include "Utils/Util.h"

using namespace RayFrame;

Mesh::Mesh()
{
  m_vertexFormat = VF_NONE;
  m_vertexBuffer = VB_NONE;
  m_indexBuffer  = IB_NONE;
  m_nIndices = 0;

  m_lastVertexCount = 0;
  m_lastVertices = NULL;
  m_lastIndices = NULL;
  m_lastFormat = NULL;
}

void convertToShorts(const unsigned *src, int nIndices, const unsigned nVertices){
  unsigned short *dest = (unsigned short *) src;

  while (nIndices > 0){
    *dest++ = *src++;
    nIndices--;
  }
}

/// Generate index
unsigned *Mesh::GetArrayIndices(const unsigned nVertices){
  unsigned *indices = new unsigned[nVertices];
  for (unsigned i = 0; i < nVertices; i++){
    indices[i] = i;
  }
  return indices;
}

/** merge all vertices and indeces into one big buffer **/
unsigned Mesh::assemble(const StreamID *aStreams, const unsigned nStreams, float **destVertices, unsigned **destIndices, bool separateArrays){
  unsigned i, j, nComp = GetComponentCount(aStreams, nStreams);
  for (i = 0; i < nStreams; i++){
    optimizeStream(aStreams[i]);
  }

  if (separateArrays){
    for (i = 0; i < nStreams; i++){
      destVertices[i] = new float[m_streams[aStreams[i]].nComponents * m_nIndices];
    }
  } else {
    *destVertices = new float[nComp * m_nIndices];
  }
  unsigned *iDest = *destIndices = new unsigned[m_nIndices];

  unsigned *iIndex = new unsigned[nStreams];
  int entryCount = m_nIndices > 8 ? m_nIndices>>3 : m_nIndices;
  Hash hash(nStreams, entryCount, m_nIndices);
  for (j = 0; j < m_nIndices; j++){
    for (i = 0; i < nStreams; i++){
      iIndex[i] = m_streams[aStreams[i]].indices[j];
    }

    unsigned index;
    if (!hash.insert(iIndex, &index)){
      if (separateArrays){
        for (i = 0; i < nStreams; i++){
          unsigned nc = m_streams[aStreams[i]].nComponents;
          float *dest = destVertices[i] + index * nc;
          memcpy(dest, m_streams[aStreams[i]].vertices + m_streams[aStreams[i]].indices[j] * nc, nc * sizeof(float));
        }
      } else {
        float *dest = *destVertices + index * nComp;
        for (i = 0; i < nStreams; i++){
          unsigned nc = m_streams[aStreams[i]].nComponents;
          memcpy(dest, m_streams[aStreams[i]].vertices + m_streams[aStreams[i]].indices[j] * nc, nc * sizeof(float));
          dest += nc;
        }
      }
    }

    *iDest++ = index;
  }
  delete iIndex;

  if (separateArrays){
    for (i = 0; i < nStreams; i++){
      destVertices[i] = (float *) realloc(destVertices[i], hash.getCount() * m_streams[aStreams[i]].nComponents * sizeof(float));
    }
  } else {
    *destVertices = (float *) realloc(*destVertices, hash.getCount() * nComp * sizeof(float));
  }
  return hash.getCount();
}

/// delete dumplicated points(points with the same coordinates
void Mesh::optimizeStream(const StreamID streamID)
{
  if (m_streams[streamID].optimized) return;

  unsigned nComp = m_streams[streamID].nComponents;
  unsigned nVert = m_streams[streamID].nVertices;

  KdTree <float> kdTree(nComp, nVert);

  unsigned *indexRemap = new unsigned[nVert];
  float *vertices = m_streams[streamID].vertices;
  for (unsigned i = 0; i < nVert; i++){
    float *vertex = vertices + i * nComp;
    unsigned index = kdTree.addUnique(vertex);

    // use kdTree to index every vertex
    indexRemap[i] = index;
    memcpy(vertices + index * nComp, vertex, nComp * sizeof(float));
  }

  unsigned *indices = m_streams[streamID].indices;
  for (unsigned j = 0; j < m_nIndices; j++){
    indices[j] = indexRemap[indices[j]];
  }

  delete indexRemap;

  // the number of vertices will be no more than before, so this realloc
  // below will not cause any trouble
  m_streams[streamID].nVertices = kdTree.getCount();
  m_streams[streamID].vertices = (float *) realloc(vertices, kdTree.getCount() * nComp * sizeof(float));
  m_streams[streamID].optimized = true;
}

BatchID Mesh::AddBatch(const unsigned startIndex, const unsigned nIndices){
  Batch batch;

  batch.startIndex = startIndex;
  batch.nIndices = nIndices;
  //  batch.startVertex = startVertex;
  //  batch.nVertices = nVertices;
  batch.startVertex = 0;
  batch.nVertices = 0;

  return m_batches.Add(batch);
}

StreamID Mesh::AddStream(const AttributeType type, const int nComponents, 
  unsigned nVertices, float *vertices, unsigned *indices, bool optimized)
{
  if (indices == NULL) indices = GetArrayIndices(nVertices);

  Stream stream;
  stream.nVertices = nVertices;
  stream.vertices  = vertices;
  stream.indices   = indices;
  stream.type      = type;
  stream.nComponents = nComponents;
  stream.optimized = optimized;

  return m_streams.Add(stream);
}

unsigned Mesh::GetVertexSize() const {
  return GetComponentCount() * sizeof(float);
}

/// Get the total number of all stream channels
unsigned Mesh::GetComponentCount() const {
  unsigned count = 0;
  for (unsigned i = 0; i < m_streams.GetCount(); i++){
    count += m_streams[i].nComponents;
  }
  return count;
}


/// Get channel number of specified stream
unsigned Mesh::GetComponentCount(const StreamID *cStreams, const unsigned nStreams) const {
  unsigned count = 0;
  for (unsigned i = 0; i < nStreams; i++){
    count += m_streams[cStreams[i]].nComponents;
  }
  return count;
}

/** Upload data to video card
  @param useCache should we keep vertex and index data in memory.
**/
unsigned Mesh::MakeDrawable(const bool useCache)
{
  if (m_streams.GetCount() == 0) return 0;

  int vertexSize = GetVertexSize();

  Renderer* renderer = OpenGLRenderer::Instance();

  if (useCache && m_lastVertices)
  { // Using cache created last time
    NotImplemented(false, "Reuse mesh buffer is not tested at all, and might cause problems.");

    if ((m_vertexFormat = renderer->AddVertexFormat(m_lastFormat, m_streams.GetCount())) == VF_NONE) 
      return 0;
    if ((m_vertexBuffer = renderer->AddVertexBuffer(m_lastVertexCount * vertexSize, STATIC, m_lastVertices)) == VB_NONE)
      return 0;

    if (m_lastVertexCount > 65535)
    {
      if ((m_indexBuffer = renderer->AddIndexBuffer(m_nIndices, 4, STATIC, m_lastIndices)) == IB_NONE) return 0;
    } 
    else 
    {
      if ((m_indexBuffer = renderer->AddIndexBuffer(m_nIndices, 2, STATIC, m_lastIndices)) == IB_NONE) return 0;
    }

    return m_lastVertexCount;
  } 
  else 
  {
    StreamID *aStreams = new StreamID[m_streams.GetCount()];

    for (unsigned i = 0; i < m_streams.GetCount(); i++)
    {
      aStreams[i] = i;
    }
    float *vertices;
    unsigned *indices;

    unsigned nVertices = assemble(aStreams, m_streams.GetCount(), &vertices, &indices, false);

    // Compute ranges for m_batches
    for (unsigned j = 0; j < m_batches.GetCount(); j++){
      unsigned minVertex = 0xFFFFFFFF;
      unsigned maxVertex = 0;

      unsigned first = m_batches[j].startIndex;
      unsigned last  = first + m_batches[j].nIndices;
      if (first < last){
        for (unsigned i = first; i < last; i++)
        {
          if (indices[i] < minVertex) minVertex = indices[i];
          if (indices[i] > maxVertex) maxVertex = indices[i];
        }

        m_batches[j].startVertex = minVertex;
        m_batches[j].nVertices = maxVertex - minVertex + 1;
      } 
      else 
      {
        // Empty batch
        m_batches[j].startVertex = 0;
        m_batches[j].nVertices = 0;
      }
    }

    FormatDesc *format = new FormatDesc[m_streams.GetCount()];
    for (unsigned i = 0; i < m_streams.GetCount(); i++)
    {
      format[i].stream = 0;
      format[i].type   = m_streams[i].type;
      format[i].format = FORMAT_FLOAT;
      format[i].size   = m_streams[i].nComponents;
    }

    if ((m_vertexFormat = renderer->AddVertexFormat(format, m_streams.GetCount())) == VF_NONE)
      return 0;
    if ((m_vertexBuffer = renderer->AddVertexBuffer(nVertices * vertexSize, STATIC, vertices)) == VB_NONE)
      return 0;

    if (nVertices > 65535)
    {
      if ((m_indexBuffer = renderer->AddIndexBuffer(m_nIndices, 4, STATIC, indices)) == IB_NONE)
        return 0;
    } 
    else 
    {
      convertToShorts(indices, m_nIndices, nVertices);

      if ((m_indexBuffer = renderer->AddIndexBuffer(m_nIndices, 2, STATIC, indices)) == IB_NONE) 
        return 0;
    }

    delete aStreams;

    if (useCache)
    {
      delete m_lastFormat;
      delete m_lastVertices;
      delete m_lastIndices;

      m_lastFormat = format;
      m_lastVertexCount = nVertices;
      m_lastVertices = vertices;
      m_lastIndices = indices;
    }
    else
    {
      delete format;
      delete vertices;
      delete indices;
    }

    renderer->CheckError();
    return nVertices;
  }
}

void Mesh::Render(double timeElapsed)
{
  assert(m_vertexBuffer != VB_NONE);
  assert(m_indexBuffer  != IB_NONE);

  Renderer* renderer = OpenGLRenderer::Instance();
  
  renderer->ChangeVertexFormat(m_vertexFormat);
  renderer->ChangeVertexBuffer(0, m_vertexBuffer);
  renderer->ChangeIndexBuffer(m_indexBuffer);
  renderer->DrawElements(PRIM_TRIANGLES, 0, m_nIndices, 0, m_lastVertexCount);
}
